#!/usr/bin/env python3
"""
离线模型管理器 - 完全不依赖网络的模型加载方案
"""

import os
import sys
import json
import yaml
import shutil
import logging
from pathlib import Path
from typing import Optional, Dict, Any

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class OfflineModelManager:
    """
    离线模型管理器
    完全不依赖 HuggingFace Hub，直接管理本地模型文件
    """
    
    def __init__(self, config_path='config/config.yaml'):
        self.config = self._load_config(config_path)
        if not self.config:
            logger.error("配置文件加载失败")
            sys.exit(1)
            
        self.model_name = self.config['model']['name']
        self.cache_dir = Path(self.config['model']['cache_dir'])
        self.model_dir = self.cache_dir / "svd_model"
        
    def _load_config(self, config_path) -> Optional[Dict]:
        """加载配置文件"""
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except Exception as e:
            logger.error(f"加载配置文件失败: {e}")
            return None
    
    def check_model_exists(self) -> bool:
        """
        检查模型文件是否存在
        不依赖 HuggingFace 的缓存结构，使用简单直接的目录结构
        """
        # 必需的模型组件
        required_components = [
            "unet",
            "vae", 
            "image_encoder",
            "feature_extractor",
            "scheduler"
        ]
        
        # 检查模型目录是否存在
        if not self.model_dir.exists():
            logger.info(f"模型目录不存在: {self.model_dir}")
            return False
        
        # 检查必需的组件
        missing_components = []
        for component in required_components:
            component_path = self.model_dir / component
            if not component_path.exists():
                missing_components.append(component)
        
        if missing_components:
            logger.info(f"缺少模型组件: {', '.join(missing_components)}")
            return False
        
        # 检查配置文件
        config_file = self.model_dir / "model_index.json"
        if not config_file.exists():
            logger.info("缺少 model_index.json")
            return False
        
        logger.info("✅ 模型文件完整")
        return True
    
    def prepare_model_structure(self):
        """
        准备模型目录结构
        创建一个清晰、简单的目录结构
        """
        logger.info("准备模型目录结构...")
        
        # 创建主目录
        self.model_dir.mkdir(parents=True, exist_ok=True)
        
        # 创建组件目录
        components = ["unet", "vae", "image_encoder", "feature_extractor", "scheduler", "tokenizer"]
        for component in components:
            (self.model_dir / component).mkdir(exist_ok=True)
        
        logger.info(f"✅ 目录结构已准备: {self.model_dir}")
        
    def create_model_index(self):
        """
        创建 model_index.json
        这是 diffusers 加载模型所必需的配置文件
        """
        model_index = {
            "_class_name": "StableVideoDiffusionPipeline",
            "_diffusers_version": "0.24.0",
            "feature_extractor": ["transformers", "CLIPImageProcessor"],
            "image_encoder": ["transformers", "CLIPVisionModelWithProjection"],
            "scheduler": ["diffusers", "EulerDiscreteScheduler"],
            "unet": ["diffusers", "UNetSpatioTemporalConditionModel"],
            "vae": ["diffusers", "AutoencoderKLTemporalDecoder"]
        }
        
        index_path = self.model_dir / "model_index.json"
        with open(index_path, 'w') as f:
            json.dump(model_index, f, indent=2)
        
        logger.info(f"✅ 创建 model_index.json: {index_path}")
    
    def provide_download_instructions(self):
        """
        提供详细的手动下载说明
        """
        print("\n" + "="*70)
        print("📥 手动下载模型文件指南")
        print("="*70)
        
        print("\n【选项 1】在可以访问网络的机器上下载")
        print("-" * 50)
        print("""
步骤 1: 在能访问 HuggingFace 的电脑上运行以下 Python 脚本：

```python
from huggingface_hub import snapshot_download
import shutil
import os

# 下载模型
model_id = "stabilityai/stable-video-diffusion-img2vid-xt"
local_dir = "./svd_download"

print("开始下载模型...")
snapshot_download(
    repo_id=model_id,
    local_dir=local_dir,
    local_dir_use_symlinks=False
)

# 打包
print("打包模型文件...")
shutil.make_archive("svd_model", 'zip', local_dir)
print("完成！请将 svd_model.zip 上传到服务器")
```

步骤 2: 将 svd_model.zip 上传到服务器的 deployment/SD Demo 目录

步骤 3: 在服务器上解压：
    unzip svd_model.zip -d models/svd_model
""")
        
        print("\n【选项 2】使用预下载的模型文件")
        print("-" * 50)
        print(f"""
如果您有预先下载的模型文件，请按以下结构放置到 {self.model_dir}：

{self.model_dir}/
├── model_index.json
├── unet/
│   ├── config.json
│   └── diffusion_pytorch_model.safetensors (或 .bin)
├── vae/
│   ├── config.json
│   └── diffusion_pytorch_model.safetensors
├── image_encoder/
│   ├── config.json
│   └── model.safetensors
├── feature_extractor/
│   └── preprocessor_config.json
└── scheduler/
    └── scheduler_config.json
""")
        
        print("\n【选项 3】从其他服务器复制")
        print("-" * 50)
        print("""
如果您有另一台已经部署成功的服务器：

在源服务器上：
    tar -czf svd_model.tar.gz models/svd_model/

复制到目标服务器：
    scp svd_model.tar.gz user@target-server:/path/to/deployment/SD Demo/

在目标服务器上解压：
    tar -xzf svd_model.tar.gz
""")
        
        print("\n" + "="*70)
        print("放置好模型文件后，重新运行部署脚本即可")
        print("="*70)
    
    def setup_offline_model(self) -> bool:
        """
        设置离线模型
        """
        logger.info("="*50)
        logger.info("离线模型设置")
        logger.info("="*50)
        
        # 检查模型是否存在
        if self.check_model_exists():
            logger.info("✅ 模型已存在，无需下载")
            return True
        
        # 准备目录结构
        self.prepare_model_structure()
        
        # 创建基本配置
        self.create_model_index()
        
        # 提供下载说明
        self.provide_download_instructions()
        
        logger.error("\n❌ 模型文件不完整，请按照上述说明手动放置模型文件")
        return False


def main():
    """主函数"""
    manager = OfflineModelManager()
    
    # 尝试设置离线模型
    success = manager.setup_offline_model()
    
    if success:
        logger.info("\n✅ 模型准备就绪，可以启动服务")
        sys.exit(0)
    else:
        logger.error("\n❌ 请先准备模型文件")
        sys.exit(1)


if __name__ == "__main__":
    main()
